#%%
%matplotlib inline
import numpy as np, matplotlib.pyplot as plt, seaborn as sns, pathlib, pandas as pd, scipy.optimize as op, lmfit as lm
%config InlineBackend.figure_format='retina'
plt.style.use('seaborn-whitegrid')
sns.set_context('paper')
sns.set_style('ticks')
cm = 1/2.54
fs = 7

def fit_func(x_data, dx_data, y_data, dy_data):
    def func(x, c): return c*x

    x_data = x_data[np.isfinite(y_data)]
    dx_data = dx_data[np.isfinite(y_data)]
    dy_data = dy_data[np.isfinite(y_data)]
    y_data = y_data[np.isfinite(y_data)]

    y_data = y_data[np.isfinite(x_data)]
    dx_data = dx_data[np.isfinite(x_data)]
    dy_data = dy_data[np.isfinite(x_data)]
    x_data = x_data[np.isfinite(x_data)]

    model = lm.Model(func)
    model.set_param_hint('c', value=15, vary=True)
    params = model.make_params()
    results = model.fit(y_data[np.isfinite(x_data)], params, weights=1/np.sqrt(dx_data**2 + dy_data**2), x=x_data[np.isfinite(x_data)])
    c = results.result.params['c'].value
    dc = results.result.params['c'].stderr
    return c, dc, func

#%%
fns = list(pathlib.Path('./').glob('velocity_data/*.csv'))
fig = plt.figure(constrained_layout=True, figsize=(18*cm,12*cm))
gs = fig.add_gridspec(2,2)
ax = fig.add_subplot(gs[0, :2])
ax2 = fig.add_subplot(gs[1, 0])
ax3 = fig.add_subplot(gs[1, 1])


fig.set_tight_layout({'rect': [0, 0, 1, 0.95], 'pad': 1.5, 'h_pad': 1.5})

df_main = pd.read_csv('velocity_data/exp_data_6_2022.txt').set_index('Unnamed: 0').rename_axis(None)


dfp = df_main[df_main['grain_kind']=='gb']
x_data = dfp.tau8_corrected.values.astype(float)
dx_data = dfp.dtau8_corrected.values.astype(float)
y_data = dfp.tau8.values.astype(float)
dy_data = dfp.dtau8.values.astype(float)
c, dc, func = fit_func(x_data, dx_data, y_data, dy_data)
x = np.linspace(0, 0.3)
ax2.plot(x, func(x, c), 'g--', label='Wall-corrected nondim. stress. Best fit: $1:%1.1f \pm %1.1f$'%(1/c, (dc/c)*(1/c)))
# plt.plot(x_data, y_data, 'go')
ax2.errorbar(x_data, y_data, xerr=dx_data, yerr=dy_data, fmt="o", c='g', lw=0.5)
print(1/c, dc, 1/(c-dc), 1/(c+dc))

# axs = axs.flatten()
count_r = -1
ustar_meas = np.zeros(8)
dustar_meas = np.zeros(8)
ustar_calc = np.zeros(8)
dustar_calc = np.zeros(8)
llims = [0.35, 0.15, 0.25]
ulims = [0.85, 0.65, 0.75]
zo = np.zeros(8)
count_t = [10, 6, 7, 9, 4, 12, 11, 5, 8]
fns_ordered = [fns[6],  fns[4], fns[5], fns[0], fns[8], fns[7], fns[1], fns[3]]
for i, fn in enumerate(fns_ordered):
    sname = str(fn).split('_piv')[0].split('/')[1]
    taustar = df_main[df_main.names==sname].tau8.values[0]
    dustar_calc[i] = df_main[df_main.names==sname].dtau8.values

    df = pd.read_csv(fn)
    y_t = df.y.values/1000
    vx_t = -df.vx_mean.values[y_t<.1]
    y_t = y_t[y_t<.1]

    ylim = y_t[np.where(vx_t <= 0.01*vx_t.max())[0][-1]]
    y_t = y_t-ylim

    if i == 0:
        ax.semilogy(vx_t+i*0.5, y_t, 'b:', label='PIV velocity-depth profile (profiles offset by 0.5 m/s for visualization)')
    else:
        ax.semilogy(vx_t+i*0.5, y_t, 'b:')

    ass = np.zeros(3)
    for jj in range(3):
        llim = llims[jj]
        y = y_t[vx_t/vx_t.max()>llim]
        vx = vx_t[vx_t/vx_t.max()>llim]
        ulim = ulims[jj]
        y = y[vx/vx.max()<ulim]
        vx = vx[vx/vx.max()<ulim]

        rhos = 2550 # sediment density (kg/m^3)
        rho = 1000 # fluid density (kg/m^3)
        k = 0.4 # von Karman constant []
        d = 0.0495 # grain diameter (m)
        g = 9.81 # gravity (m/s^2)
        nu = 1e-6 # kinematic viscosity (m^2/s)
        taub = taustar * ( (rhos-rho)*g*d )
        ustar = np.sqrt(taub/rho)
        Restar = ustar*d/nu
        z0 = d/30; # vertical coordinate at which law of wall velocity = 0 (m)

        vx_LOW = ustar/k * np.log(y[y > 0]/z0) # law of the wall velocity (m/s)

        # ax.text(0.1+i*0.55, 0.1, r'$\tau^* = %0.3f$' % taustar)

        if jj == 2:
            if i == 0:
                ax.semilogy(vx+i*0.5, y, 'b-', lw=2, label='Fitting region for Log law of the wall')
            else:
                ax.semilogy(vx+i*0.5, y, 'b-', lw=2)

        func = lambda x, a, b: np.exp(a*x + b)
        model = lm.Model(func)
        model.set_param_hint('a', value=1, vary=True)
        model.set_param_hint('b', value=1, vary=True)
        params = model.make_params()
        results = model.fit(y, params, x=vx)
        a = results.result.params['a'].value
        ass[jj] = a
        da = results.result.params['a'].stderr
        # pa = da/a
        # print(pa)
        b = results.result.params['b'].value
        db = results.result.params['b'].stderr

    a = ass.mean()
    da = ass.std()
    pa = da/a
    # popt, pcov = op.curve_fit(func, vx, y)
    # print(popt, a, b)

    x = np.linspace(0.1, 1.5)
    ax.plot(x+i*0.5, func(x, a, b), 'k-', lw=0.8)

    ustar_meas[i] = (k/a)**2  / ( (rhos/rho-1)*g*d )
    dustar_meas[i] = ustar_meas[i] * pa
    # print(ustar_meas[i], dustar_meas[i])
    ustar_calc[i] = ustar**2 / ( (rhos/rho-1)*g*d )
    # print(ustar_meas[i])
    zo[i] = b
    ax.set_ylim(1e-3, .2)

    ax.set_xlabel('Downstream velocity (m/s)', fontsize=fs)
    ax.set_ylabel('Elevation above bed (m)', fontsize=fs)
    ax.legend(loc=4, fontsize=7, fancybox=True, frameon=True, framealpha=1)

# fig = plt.figure(figsize=(4,3))
ax2.errorbar(ustar_meas, ustar_calc, xerr=dustar_meas, yerr=dustar_calc, fmt="o", c='b', lw=0.5)
# print(ustar_meas, dustar_meas, ustar_calc, dustar_calc)
c, dc, func = fit_func(ustar_meas, dustar_meas, ustar_calc, dustar_calc)
print(1/c)
x = np.linspace(0, .3)
ax2.plot(x, func(c, x), 'b--', label='PIV-derived nondim. stress. $1:%1.1f \pm %1.1f$'%(1/c, (dc/c)*(1/c)))
# popt, pcov = op.curve_fit(func, ustar_meas, ustar_calc)
# ax2.plot(x, func(x, *popt), 'k--', label='1:%0.1f'%(1/popt))
# ax2.plot(x, (0.019/0.05)*x, 'k:', label='1:%0.1f'%(0.05/0.019))
# ax2.plot(x, x, 'g--', label='1:1')
ax2.legend(loc=2, fontsize=6, fancybox=True, frameon=True, framealpha=1)
sns.despine()
ax2.set_xlabel('PIV-derived or wall-corrected nondimensional shear stress')
ax2.set_ylabel(r'$\tau^* = \frac{\rho g R S}{(\rho_s-\rho) g d_o}$')
ax2.set_ylim(0, .15)
ax2.set_xlim(-0.03, .4)
plt.tight_layout()
sns.despine()

# natural grains
dfp = df_main[df_main['grain_kind']=='ng']
x_data = dfp.tau8_corrected.values.astype(float)
dx_data = dfp.dtau8_corrected.values.astype(float)
y_data = dfp.tau8.values.astype(float)
dy_data = dfp.dtau8.values.astype(float)
c, dc, func = fit_func(x_data, dx_data, y_data, dy_data)
x = np.linspace(0, 0.3)
ax3.plot(x, func(x, c), 'g--')
# plt.plot(x_data, y_data, 'go')
ax3.errorbar(x_data, y_data, xerr=dx_data, yerr=dy_data, fmt="o", c='g', lw=0.5, label='Best fit: $1:%1.1f \pm %1.1f$'%(1/c, (dc/c)*(1/c)))
ax3.set_xlabel('Wall-corrected nondimensional shear stress')
ax3.set_ylim(0, .15)
ax3.set_xlim(-0.03, .4)
ax3.legend(loc=2, fontsize=7, fancybox=True, frameon=True, framealpha=1)

ax.text(5.35, 0.1, 'a', fontsize=fs+1, weight='bold')
ax2.text(0.37, 0.13, 'b', fontsize=fs+1, weight='bold')
ax2.text(0.25, 0.025, 'Spheres', fontsize=fs)
ax3.text(0.37, 0.13, 'c', fontsize=fs+1, weight='bold')
ax3.text(0.25, 0.025, 'Natural Gravel 1', fontsize=fs)


print(1/c, dc, 1/(c-dc), 1/(c+dc))
plt.savefig('Deal_EDfig9.jpg', dpi=300)
# %%
